# server.py
import asyncio
import websockets
from fastapi import FastAPI, WebSocket
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import uvicorn
import json

app = FastAPI()

# 静态文件服务
app.mount("/static", StaticFiles(directory="static"), name="static")

# WebSocket连接管理器
class ConnectionManager:
    def __init__(self):
        self.active_connections = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, data: dict):
        for connection in self.active_connections:
            try:
                await connection.send_json(data)
            except Exception as e:
                print(f"Error sending data: {e}")
                self.disconnect(connection)

manager = ConnectionManager()

# 模拟数据生成任务
async def data_generator():
    import random
    import time
    while True:
        # 这里应该替换为从C++应用获取真实数据
        # 模拟1000Hz采样率，每次发送10个点(10ms数据)
        data = {
            "type": "waveform",
            "data": [random.randint(-32768, 32767) for _ in range(10)],
            "timestamp": time.time()
        }
        await manager.broadcast(data)
        await asyncio.sleep(0.01)  # 10ms间隔

# WebSocket端点
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            # 保持连接开放
            await websocket.receive_text()
    except Exception as e:
        print(f"WebSocket error: {e}")
    finally:
        manager.disconnect(websocket)

# 提供HTML页面
@app.get("/")
async def get():
    with open("static/index.html") as f:
        return HTMLResponse(f.read())

async def main():
    # 启动数据生成任务
    asyncio.create_task(data_generator())
    
    # 启动Web服务器
    config = uvicorn.Config(app, host="0.0.0.0", port=8000)
    server = uvicorn.Server(config)
    await server.serve()

if __name__ == "__main__":
    asyncio.run(main())